import logging
import os
import pickle
import numpy as np
import pandas as pd
from typing import List, Dict, Any, Optional, Tuple, Callable, Union, Type
from datetime import datetime
import matplotlib.pyplot as plt
import statsmodels.api as sm
from scipy.stats import spearmanr, pearsonr, ttest_ind
from sklearn.linear_model import LinearRegression, Ridge, Lasso
from sklearn.preprocessing import StandardScaler
import seaborn as sns

class Factor:
    """单因子分析类，用于创建和分析单个因子"""
    
    def __init__(self, name: str, func: Callable = None, data: pd.Series = None, 
               description: str = None, category: str = None):
        """
        初始化因子
        
        参数:
            name: 因子名称
            func: 因子计算函数，接受pd.DataFrame，返回pd.Series
            data: 预计算的因子数据
            description: 因子描述
            category: 因子类别
        """
        self.name = name
        self.func = func
        self.data = data
        self.description = description or f"因子: {name}"
        self.category = category or "未分类"
        self.ic_series = None
        self.stats = {}
        
    def calculate(self, data: pd.DataFrame) -> pd.Series:
        """
        计算因子值
        
        参数:
            data: 输入数据
            
        返回:
            pd.Series: 因子值
        """
        if self.func is not None:
            try:
                result = self.func(data)
                if not isinstance(result, pd.Series):
                    result = pd.Series(result, index=data.index)
                self.data = result
                return result
            except Exception as e:
                logging.error(f"计算因子 {self.name} 时出错: {str(e)}")
                raise
        elif self.data is not None:
            return self.data
        else:
            raise ValueError("未指定因子计算函数或预计算数据")
            
    def rank(self) -> pd.Series:
        """
        计算因子值排名
        
        返回:
            pd.Series: 因子排名
        """
        if self.data is None:
            raise ValueError("因子数据未计算")
            
        return self.data.rank(method='average')
        
    def normalize(self, method: str = 'zscore') -> pd.Series:
        """
        标准化因子值
        
        参数:
            method: 标准化方法，可选 'zscore', 'minmax', 'rank'
            
        返回:
            pd.Series: 标准化后的因子值
        """
        if self.data is None:
            raise ValueError("因子数据未计算")
            
        if method == 'zscore':
            # Z-score标准化
            return (self.data - self.data.mean()) / self.data.std()
        elif method == 'minmax':
            # Min-Max标准化
            return (self.data - self.data.min()) / (self.data.max() - self.data.min())
        elif method == 'rank':
            # 排名标准化
            return self.rank() / len(self.data)
        else:
            raise ValueError(f"未知的标准化方法: {method}")
            
    def winsorize(self, limits: Tuple[float, float] = (0.01, 0.01)) -> pd.Series:
        """
        缩尾处理
        
        参数:
            limits: 缩尾范围，例如(0.01, 0.01)表示缩尾1%
            
        返回:
            pd.Series: 缩尾处理后的因子值
        """
        if self.data is None:
            raise ValueError("因子数据未计算")
            
        lower_limit = self.data.quantile(limits[0])
        upper_limit = self.data.quantile(1 - limits[1])
        
        winsorized = self.data.copy()
        winsorized[winsorized < lower_limit] = lower_limit
        winsorized[winsorized > upper_limit] = upper_limit
        
        return winsorized
        
    def neutralize(self, to_neutralize: List[pd.Series]) -> pd.Series:
        """
        因子中性化处理
        
        参数:
            to_neutralize: 用于中性化的变量列表（如行业、市值等）
            
        返回:
            pd.Series: 中性化后的因子值
        """
        if self.data is None:
            raise ValueError("因子数据未计算")
            
        # 构建数据框
        df = pd.DataFrame({'y': self.data})
        
        # 添加中性化变量
        for i, var in enumerate(to_neutralize):
            if isinstance(var, pd.Series):
                df[f'x{i}'] = var
            else:
                raise ValueError(f"中性化变量必须是pd.Series类型，第{i}个变量类型错误")
                
        # 删除缺失值
        df = df.dropna()
        
        # 线性回归
        X = df.drop('y', axis=1)
        y = df['y']
        
        if len(X) == 0:
            logging.warning(f"中性化变量与因子 {self.name} 无共同数据点")
            return self.data
            
        model = LinearRegression()
        model.fit(X, y)
        
        # 残差是中性化后的因子
        residuals = y - model.predict(X)
        
        # 将残差扩展到原始索引，未回归的值保持不变
        neutralized = self.data.copy()
        neutralized.loc[residuals.index] = residuals
        
        return neutralized
        
    def get_groups(self, n_groups: int = 5) -> pd.Series:
        """
        获取因子分组
        
        参数:
            n_groups: 分组数
            
        返回:
            pd.Series: 分组标签
        """
        if self.data is None:
            raise ValueError("因子数据未计算")
            
        # 使用qcut进行分组
        return pd.qcut(self.data, n_groups, labels=list(range(1, n_groups + 1)))
        
    def calc_ic(self, returns: pd.Series, method: str = 'spearman', 
               periods: List[int] = [1, 5, 10, 20], group_neutral: bool = False) -> Dict[str, Any]:
        """
        计算因子IC值
        
        参数:
            returns: 收益率数据
            method: 相关系数方法，可选 'spearman', 'pearson'
            periods: 计算多期IC的周期列表
            group_neutral: 是否进行分组中性化
            
        返回:
            Dict: IC统计信息
        """
        if self.data is None:
            raise ValueError("因子数据未计算")
            
        # 确保因子与收益率有相同的索引
        common_index = self.data.index.intersection(returns.index)
        if len(common_index) == 0:
            raise ValueError("因子数据和收益率数据无共同索引")
            
        factor = self.data.loc[common_index]
        ret = returns.loc[common_index]
        
        ic_dict = {}
        ic_series_dict = {}
        
        for period in periods:
            # 计算未来N期收益率
            future_returns = ret.shift(-period)
            
            # 删除缺失值
            mask = ~(factor.isna() | future_returns.isna())
            clean_factor = factor[mask]
            clean_returns = future_returns[mask]
            
            if len(clean_factor) == 0:
                logging.warning(f"因子 {self.name} 与 {period} 期收益率无共同有效数据")
                continue
                
            # 计算IC
            if method == 'spearman':
                ic, p_value = spearmanr(clean_factor, clean_returns)
            elif method == 'pearson':
                ic, p_value = pearsonr(clean_factor, clean_returns)
            else:
                raise ValueError(f"未知的相关系数方法: {method}")
                
            # 记录IC值
            ic_dict[f'ic_{period}'] = ic
            ic_dict[f'ic_{period}_p_value'] = p_value
            
            # 如果是多期数据，计算滚动IC
            if len(clean_factor) > 10:  # 至少需要10个数据点
                # 计算滚动IC
                factor_shifts = pd.DataFrame({
                    'factor': clean_factor,
                    'returns': clean_returns
                })
                
                # 计算每个时间点的IC
                if method == 'spearman':
                    rolling_ic = factor_shifts.rolling(window=min(20, len(factor_shifts))).apply(
                        lambda x: spearmanr(x['factor'], x['returns'])[0] if len(x) > 2 else np.nan
                    )
                else:
                    rolling_ic = factor_shifts.rolling(window=min(20, len(factor_shifts))).apply(
                        lambda x: pearsonr(x['factor'], x['returns'])[0] if len(x) > 2 else np.nan
                    )
                    
                ic_series_dict[f'ic_{period}'] = rolling_ic
                
                # 计算IR值(IC均值/IC标准差)
                ic_mean = rolling_ic.mean()
                ic_std = rolling_ic.std()
                ir = ic_mean / ic_std if ic_std > 0 else 0
                
                ic_dict[f'ir_{period}'] = ir
        
        # 保存IC系列，用于后续分析
        self.ic_series = ic_series_dict
        
        # 保存统计结果
        self.stats.update(ic_dict)
        
        return ic_dict
        
    def backtest(self, returns: pd.Series, n_groups: int = 5, 
                fee_rate: float = 0.0015, turnover_rate: float = 0.2) -> Dict[str, Any]:
        """
        进行分组回测
        
        参数:
            returns: 收益率数据
            n_groups: 分组数
            fee_rate: 交易费率
            turnover_rate: 换手率
            
        返回:
            Dict: 回测结果
        """
        if self.data is None:
            raise ValueError("因子数据未计算")
            
        # 获取因子分组
        groups = self.get_groups(n_groups=n_groups)
        
        # 确保因子与收益率有相同的索引
        common_index = groups.index.intersection(returns.index)
        if len(common_index) == 0:
            raise ValueError("因子分组和收益率数据无共同索引")
            
        groups = groups.loc[common_index]
        ret = returns.loc[common_index]
        
        # 分组收益率
        group_returns = {}
        for group in range(1, n_groups + 1):
            group_mask = (groups == group)
            if group_mask.sum() > 0:
                group_returns[group] = ret[group_mask]
                
        # 多空组合收益率（最高分组减最低分组）
        if 1 in group_returns and n_groups in group_returns:
            long_short_returns = group_returns[n_groups] - group_returns[1]
            
            # 计算交易成本
            cost = fee_rate * turnover_rate
            long_short_returns_after_cost = long_short_returns - cost
            
            # 计算累积收益
            cum_returns = {}
            cum_returns_after_cost = {}
            
            for group, group_ret in group_returns.items():
                cum_returns[group] = (1 + group_ret).cumprod()
                
            cum_long_short = (1 + long_short_returns).cumprod()
            cum_long_short_after_cost = (1 + long_short_returns_after_cost).cumprod()
            
            # 计算年化收益率
            years = (ret.index[-1] - ret.index[0]).days / 365.25
            annual_returns = {}
            
            for group, cum_ret in cum_returns.items():
                annual_returns[group] = (cum_ret.iloc[-1] ** (1 / years)) - 1
                
            annual_long_short = (cum_long_short.iloc[-1] ** (1 / years)) - 1
            annual_long_short_after_cost = (cum_long_short_after_cost.iloc[-1] ** (1 / years)) - 1
            
            # 计算夏普比率
            sharpe_ratios = {}
            
            for group, group_ret in group_returns.items():
                annual_vol = group_ret.std() * np.sqrt(252)
                sharpe_ratios[group] = annual_returns[group] / annual_vol if annual_vol > 0 else 0
                
            annual_long_short_vol = long_short_returns.std() * np.sqrt(252)
            sharpe_long_short = annual_long_short / annual_long_short_vol if annual_long_short_vol > 0 else 0
            
            annual_long_short_after_cost_vol = long_short_returns_after_cost.std() * np.sqrt(252)
            sharpe_long_short_after_cost = annual_long_short_after_cost / annual_long_short_after_cost_vol if annual_long_short_after_cost_vol > 0 else 0
            
            # 计算最大回撤
            drawdowns = {}
            
            for group, cum_ret in cum_returns.items():
                rolling_max = cum_ret.cummax()
                drawdown = (cum_ret - rolling_max) / rolling_max
                drawdowns[group] = drawdown.min()
                
            rolling_max_ls = cum_long_short.cummax()
            drawdown_ls = (cum_long_short - rolling_max_ls) / rolling_max_ls
            max_drawdown_ls = drawdown_ls.min()
            
            rolling_max_ls_ac = cum_long_short_after_cost.cummax()
            drawdown_ls_ac = (cum_long_short_after_cost - rolling_max_ls_ac) / rolling_max_ls_ac
            max_drawdown_ls_ac = drawdown_ls_ac.min()
            
            # 汇总结果
            result = {
                'group_returns': group_returns,
                'long_short_returns': long_short_returns,
                'long_short_returns_after_cost': long_short_returns_after_cost,
                'cum_returns': cum_returns,
                'cum_long_short': cum_long_short,
                'cum_long_short_after_cost': cum_long_short_after_cost,
                'annual_returns': annual_returns,
                'annual_long_short': annual_long_short,
                'annual_long_short_after_cost': annual_long_short_after_cost,
                'sharpe_ratios': sharpe_ratios,
                'sharpe_long_short': sharpe_long_short,
                'sharpe_long_short_after_cost': sharpe_long_short_after_cost,
                'max_drawdowns': drawdowns,
                'max_drawdown_long_short': max_drawdown_ls,
                'max_drawdown_long_short_after_cost': max_drawdown_ls_ac
            }
            
            # 保存回测统计
            bt_stats = {
                'annual_long_short': annual_long_short,
                'annual_long_short_after_cost': annual_long_short_after_cost,
                'sharpe_long_short': sharpe_long_short,
                'sharpe_long_short_after_cost': sharpe_long_short_after_cost,
                'max_drawdown_long_short': max_drawdown_ls,
                'max_drawdown_long_short_after_cost': max_drawdown_ls_ac
            }
            self.stats.update(bt_stats)
            
            return result
        else:
            logging.warning(f"因子 {self.name} 分组不完整，无法计算多空组合收益")
            return {'group_returns': group_returns}
            
    def plot(self, returns: pd.Series = None, n_groups: int = 5, figsize: Tuple[int, int] = (15, 10)) -> plt.Figure:
        """
        绘制因子分析图表
        
        参数:
            returns: 收益率数据
            n_groups: 分组数
            figsize: 图表大小
            
        返回:
            plt.Figure: 图表对象
        """
        if self.data is None:
            raise ValueError("因子数据未计算")
            
        fig = plt.figure(figsize=figsize)
        
        # 布局设置
        if returns is not None:
            grid = plt.GridSpec(3, 2, figure=fig)
            ax1 = fig.add_subplot(grid[0, 0])  # 因子分布
            ax2 = fig.add_subplot(grid[0, 1])  # 因子IC分布
            ax3 = fig.add_subplot(grid[1, :])  # 分组收益率
            ax4 = fig.add_subplot(grid[2, :])  # 多空组合累积收益
        else:
            grid = plt.GridSpec(1, 2, figure=fig)
            ax1 = fig.add_subplot(grid[0, 0])  # 因子分布
            ax2 = fig.add_subplot(grid[0, 1])  # 因子自相关
            ax3 = ax4 = None
            
        # 1. 因子分布
        sns.histplot(self.data.dropna(), ax=ax1, kde=True)
        ax1.set_title(f"因子分布: {self.name}")
        ax1.set_xlabel("因子值")
        ax1.set_ylabel("频率")
        
        # 2. 因子自相关或IC分布
        if self.ic_series is not None and len(self.ic_series) > 0:
            # 如果有IC系列，绘制IC分布
            for period, ic in self.ic_series.items():
                if len(ic.dropna()) > 0:
                    sns.histplot(ic.dropna(), ax=ax2, kde=True, label=period)
            ax2.set_title(f"因子IC分布: {self.name}")
            ax2.set_xlabel("IC值")
            ax2.set_ylabel("频率")
            ax2.legend()
        else:
            # 否则绘制因子自相关
            lags = min(20, len(self.data) // 10)
            if lags > 1:
                pd.plotting.autocorrelation_plot(self.data.dropna(), ax=ax2)
                ax2.set_title(f"因子自相关: {self.name}")
                ax2.set_xlim(0, lags)
        
        # 如果提供了收益率数据，绘制回测相关图表
        if returns is not None and ax3 is not None and ax4 is not None:
            backtest_result = self.backtest(returns, n_groups=n_groups)
            
            # 3. 分组收益率
            if 'annual_returns' in backtest_result:
                groups = list(backtest_result['annual_returns'].keys())
                returns_values = [backtest_result['annual_returns'][g] for g in groups]
                ax3.bar(groups, returns_values)
                ax3.set_title(f"各分组年化收益率: {self.name}")
                ax3.set_xlabel("分组")
                ax3.set_ylabel("年化收益率")
                
                # 添加数值标签
                for i, v in enumerate(returns_values):
                    ax3.text(i + 1, v, f"{v:.2%}", ha='center', va='bottom')
            
            # 4. 多空组合累积收益
            if 'cum_long_short' in backtest_result and 'cum_long_short_after_cost' in backtest_result:
                backtest_result['cum_long_short'].plot(ax=ax4, label='多空组合')
                backtest_result['cum_long_short_after_cost'].plot(ax=ax4, label='多空组合(扣费)')
                
                # 添加各分组累积收益
                if 'cum_returns' in backtest_result:
                    for group, cum_ret in backtest_result['cum_returns'].items():
                        cum_ret.plot(ax=ax4, label=f'分组{group}', alpha=0.5, linestyle='--')
                        
                ax4.set_title(f"累积收益曲线: {self.name}")
                ax4.set_xlabel("日期")
                ax4.set_ylabel("累积收益")
                ax4.legend()
                ax4.grid(True)
        
        plt.tight_layout()
        return fig


class FactorLibrary:
    """因子库类，用于管理和分析多个因子"""
    
    def __init__(self):
        """初始化因子库"""
        self.factors = {}
        self.category_map = {}
        
    def add_factor(self, factor: Factor) -> 'FactorLibrary':
        """
        添加因子
        
        参数:
            factor: 因子对象
            
        返回:
            self: 支持链式调用
        """
        self.factors[factor.name] = factor
        
        # 更新类别映射
        if factor.category not in self.category_map:
            self.category_map[factor.category] = []
        self.category_map[factor.category].append(factor.name)
        
        return self
        
    def get_factor(self, name: str) -> Optional[Factor]:
        """
        获取因子
        
        参数:
            name: 因子名称
            
        返回:
            Factor: 因子对象
        """
        return self.factors.get(name)
        
    def get_factors_by_category(self, category: str) -> List[Factor]:
        """
        按类别获取因子
        
        参数:
            category: 因子类别
            
        返回:
            List[Factor]: 因子对象列表
        """
        if category not in self.category_map:
            return []
            
        return [self.factors[name] for name in self.category_map[category]]
        
    def list_factors(self) -> pd.DataFrame:
        """
        列出所有因子
        
        返回:
            pd.DataFrame: 因子信息表
        """
        data = []
        for name, factor in self.factors.items():
            data.append({
                'name': name,
                'category': factor.category,
                'description': factor.description,
                'has_data': factor.data is not None
            })
            
        return pd.DataFrame(data)
        
    def list_categories(self) -> List[str]:
        """
        列出所有因子类别
        
        返回:
            List[str]: 类别列表
        """
        return list(self.category_map.keys())
        
    def calculate_factors(self, data: pd.DataFrame, names: List[str] = None) -> Dict[str, pd.Series]:
        """
        计算多个因子值
        
        参数:
            data: 输入数据
            names: 指定要计算的因子名称列表
            
        返回:
            Dict[str, pd.Series]: 因子名称到因子值的映射
        """
        if names is None:
            names = list(self.factors.keys())
            
        results = {}
        for name in names:
            if name in self.factors:
                try:
                    factor_data = self.factors[name].calculate(data)
                    results[name] = factor_data
                except Exception as e:
                    logging.error(f"计算因子 {name} 时出错: {str(e)}")
            else:
                logging.warning(f"因子 {name} 不存在")
                
        return results
        
    def get_factor_matrix(self, names: List[str] = None) -> pd.DataFrame:
        """
        获取因子矩阵
        
        参数:
            names: 指定要包含的因子名称列表
            
        返回:
            pd.DataFrame: 因子矩阵，列为因子，行为观测
        """
        if names is None:
            names = [name for name, factor in self.factors.items() if factor.data is not None]
            
        data = {}
        common_index = None
        
        for name in names:
            factor = self.get_factor(name)
            if factor is not None and factor.data is not None:
                data[name] = factor.data
                if common_index is None:
                    common_index = factor.data.index
                else:
                    common_index = common_index.intersection(factor.data.index)
                    
        if not data:
            return pd.DataFrame()
            
        # 对齐索引
        aligned_data = {}
        for name, series in data.items():
            aligned_data[name] = series.loc[common_index]
            
        return pd.DataFrame(aligned_data)
        
    def correlation_matrix(self, names: List[str] = None, method: str = 'pearson') -> pd.DataFrame:
        """
        计算因子相关性矩阵
        
        参数:
            names: 指定要包含的因子名称列表
            method: 相关系数方法
            
        返回:
            pd.DataFrame: 相关性矩阵
        """
        factor_matrix = self.get_factor_matrix(names)
        if factor_matrix.empty:
            return pd.DataFrame()
            
        return factor_matrix.corr(method=method)
        
    def factor_ic_summary(self, returns: pd.Series, names: List[str] = None, 
                        method: str = 'spearman') -> pd.DataFrame:
        """
        计算多个因子的IC汇总
        
        参数:
            returns: 收益率数据
            names: 指定要包含的因子名称列表
            method: 相关系数方法
            
        返回:
            pd.DataFrame: IC汇总
        """
        if names is None:
            names = [name for name, factor in self.factors.items() if factor.data is not None]
            
        result = []
        for name in names:
            factor = self.get_factor(name)
            if factor is not None and factor.data is not None:
                ic_data = factor.calc_ic(returns, method=method)
                ic_data['name'] = name
                ic_data['category'] = factor.category
                result.append(ic_data)
                
        if not result:
            return pd.DataFrame()
            
        return pd.DataFrame(result)
        
    def backtest_summary(self, returns: pd.Series, names: List[str] = None, 
                       n_groups: int = 5) -> pd.DataFrame:
        """
        计算多个因子的回测汇总
        
        参数:
            returns: 收益率数据
            names: 指定要包含的因子名称列表
            n_groups: 分组数
            
        返回:
            pd.DataFrame: 回测汇总
        """
        if names is None:
            names = [name for name, factor in self.factors.items() if factor.data is not None]
            
        result = []
        for name in names:
            factor = self.get_factor(name)
            if factor is not None and factor.data is not None:
                try:
                    bt_result = factor.backtest(returns, n_groups=n_groups)
                    summary = {
                        'name': name,
                        'category': factor.category
                    }
                    
                    # 提取关键回测指标
                    for key in ['annual_long_short', 'annual_long_short_after_cost', 
                               'sharpe_long_short', 'sharpe_long_short_after_cost', 
                               'max_drawdown_long_short', 'max_drawdown_long_short_after_cost']:
                        if key in factor.stats:
                            summary[key] = factor.stats[key]
                            
                    result.append(summary)
                except Exception as e:
                    logging.error(f"因子 {name} 回测失败: {str(e)}")
                    
        if not result:
            return pd.DataFrame()
            
        return pd.DataFrame(result)
        
    def save(self, path: str) -> None:
        """
        保存因子库
        
        参数:
            path: 保存路径
        """
        os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
        
        with open(path, 'wb') as f:
            pickle.dump({
                'factors': self.factors,
                'category_map': self.category_map
            }, f)
            
        logging.info(f"因子库已保存到: {path}")
        
    @classmethod
    def load(cls, path: str) -> 'FactorLibrary':
        """
        加载因子库
        
        参数:
            path: 因子库路径
            
        返回:
            FactorLibrary: 因子库对象
        """
        with open(path, 'rb') as f:
            data = pickle.load(f)
            
        library = cls()
        library.factors = data['factors']
        library.category_map = data['category_map']
        
        return library
        
    def plot_correlation_matrix(self, names: List[str] = None, method: str = 'pearson', 
                              figsize: Tuple[int, int] = (10, 8)) -> plt.Figure:
        """
        绘制因子相关性矩阵热力图
        
        参数:
            names: 指定要包含的因子名称列表
            method: 相关系数方法
            figsize: 图表大小
            
        返回:
            plt.Figure: 图表对象
        """
        corr_matrix = self.correlation_matrix(names, method)
        if corr_matrix.empty:
            raise ValueError("没有有效的因子数据计算相关性矩阵")
            
        fig, ax = plt.subplots(figsize=figsize)
        mask = np.triu(np.ones_like(corr_matrix, dtype=bool))
        cmap = sns.diverging_palette(220, 10, as_cmap=True)
        
        sns.heatmap(corr_matrix, mask=mask, cmap=cmap, vmax=1, vmin=-1, center=0,
                   annot=True, fmt=".2f", square=True, linewidths=.5, ax=ax)
                   
        ax.set_title(f"因子相关性矩阵 ({method})")
        plt.tight_layout()
        
        return fig
        
    def plot_ic_comparison(self, returns: pd.Series, names: List[str] = None, 
                         method: str = 'spearman', period: int = 1,
                         figsize: Tuple[int, int] = (12, 6)) -> plt.Figure:
        """
        绘制多因子IC对比图
        
        参数:
            returns: 收益率数据
            names: 指定要包含的因子名称列表
            method: 相关系数方法
            period: IC周期
            figsize: 图表大小
            
        返回:
            plt.Figure: 图表对象
        """
        # 计算IC
        for name, factor in self.factors.items():
            if factor.data is not None and (names is None or name in names):
                factor.calc_ic(returns, method=method)
                
        # 提取IC值
        ic_values = {}
        for name, factor in self.factors.items():
            if factor.stats and f'ic_{period}' in factor.stats and (names is None or name in names):
                ic_values[name] = factor.stats[f'ic_{period}']
                
        if not ic_values:
            raise ValueError("没有有效的IC值")
            
        # 绘制IC对比柱状图
        fig, ax = plt.subplots(figsize=figsize)
        
        # 按IC值排序
        sorted_items = sorted(ic_values.items(), key=lambda x: abs(x[1]), reverse=True)
        factors = [item[0] for item in sorted_items]
        ic_vals = [item[1] for item in sorted_items]
        
        # 根据IC值的正负设置颜色
        colors = ['g' if ic >= 0 else 'r' for ic in ic_vals]
        
        bars = ax.bar(factors, ic_vals, color=colors)
        ax.set_title(f"{period}期IC值对比 ({method})")
        ax.set_xlabel("因子")
        ax.set_ylabel("IC值")
        ax.axhline(y=0, color='black', linestyle='-', alpha=0.3)
        
        # 添加数值标签
        for bar in bars:
            height = bar.get_height()
            ax.text(bar.get_x() + bar.get_width()/2., 
                   height if height > 0 else height - 0.02,
                   f"{height:.3f}",
                   ha='center', va='bottom' if height > 0 else 'top')
                   
        plt.xticks(rotation=45, ha='right')
        plt.tight_layout()
        
        return fig
        
    def plot_returns_comparison(self, returns: pd.Series, names: List[str] = None, 
                              n_groups: int = 5, figsize: Tuple[int, int] = (12, 8)) -> plt.Figure:
        """
        绘制多因子回测收益对比图
        
        参数:
            returns: 收益率数据
            names: 指定要包含的因子名称列表
            n_groups: 分组数
            figsize: 图表大小
            
        返回:
            plt.Figure: 图表对象
        """
        # 回测所有因子
        for name, factor in self.factors.items():
            if factor.data is not None and (names is None or name in names):
                try:
                    factor.backtest(returns, n_groups=n_groups)
                except Exception as e:
                    logging.error(f"因子 {name} 回测失败: {str(e)}")
                    
        # 收集多空组合累积收益
        cum_returns = {}
        annual_returns = {}
        
        for name, factor in self.factors.items():
            if 'annual_long_short' in factor.stats and (names is None or name in names):
                annual_returns[name] = factor.stats['annual_long_short']
                
        if not annual_returns:
            raise ValueError("没有有效的回测结果")
            
        # 创建图表
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=figsize)
        
        # 1. 年化收益率对比
        sorted_items = sorted(annual_returns.items(), key=lambda x: x[1], reverse=True)
        factors = [item[0] for item in sorted_items]
        returns_vals = [item[1] for item in sorted_items]
        
        bars = ax1.bar(factors, returns_vals, color='steelblue')
        ax1.set_title("多空组合年化收益率对比")
        ax1.set_xlabel("因子")
        ax1.set_ylabel("年化收益率")
        ax1.axhline(y=0, color='black', linestyle='-', alpha=0.3)
        
        # 添加数值标签
        for bar in bars:
            height = bar.get_height()
            ax1.text(bar.get_x() + bar.get_width()/2., 
                    height if height > 0 else height - 0.02,
                    f"{height:.2%}",
                    ha='center', va='bottom' if height > 0 else 'top')
                    
        ax1.set_xticklabels(factors, rotation=45, ha='right')
        
        # 2. 夏普比率对比
        sharpe_ratios = {}
        for name, factor in self.factors.items():
            if 'sharpe_long_short' in factor.stats and (names is None or name in names):
                sharpe_ratios[name] = factor.stats['sharpe_long_short']
                
        if sharpe_ratios:
            sorted_items = sorted(sharpe_ratios.items(), key=lambda x: x[1], reverse=True)
            factors = [item[0] for item in sorted_items]
            sharpe_vals = [item[1] for item in sorted_items]
            
            bars = ax2.bar(factors, sharpe_vals, color='coral')
            ax2.set_title("多空组合夏普比率对比")
            ax2.set_xlabel("因子")
            ax2.set_ylabel("夏普比率")
            ax2.axhline(y=0, color='black', linestyle='-', alpha=0.3)
            
            # 添加数值标签
            for bar in bars:
                height = bar.get_height()
                ax2.text(bar.get_x() + bar.get_width()/2., 
                        height if height > 0 else height - 0.02,
                        f"{height:.2f}",
                        ha='center', va='bottom' if height > 0 else 'top')
                        
            ax2.set_xticklabels(factors, rotation=45, ha='right')
        
        plt.tight_layout()
        
        return fig


class MultiFactorModel:
    """多因子模型类，用于构建和测试多因子组合"""
    
    def __init__(self, name: str = "多因子模型"):
        """
        初始化多因子模型
        
        参数:
            name: 模型名称
        """
        self.name = name
        self.factors = {}  # 因子字典
        self.weights = {}  # 因子权重
        self.factor_data = None  # 因子数据矩阵
        self.model = None  # 回归模型
        self.combined_factor = None  # 组合因子
        
    def add_factor(self, factor: Factor, weight: float = None) -> 'MultiFactorModel':
        """
        添加因子
        
        参数:
            factor: 因子对象
            weight: 因子权重
            
        返回:
            self: 支持链式调用
        """
        if factor.data is None:
            raise ValueError(f"因子 {factor.name} 数据未计算")
            
        self.factors[factor.name] = factor
        
        if weight is not None:
            self.weights[factor.name] = weight
            
        return self
        
    def set_weight(self, name: str, weight: float) -> 'MultiFactorModel':
        """
        设置因子权重
        
        参数:
            name: 因子名称
            weight: 因子权重
            
        返回:
            self: 支持链式调用
        """
        if name not in self.factors:
            raise ValueError(f"因子 {name} 不存在")
            
        self.weights[name] = weight
        
        return self
        
    def normalize_weights(self) -> 'MultiFactorModel':
        """
        标准化因子权重（使权重和为1）
        
        返回:
            self: 支持链式调用
        """
        total_weight = sum(self.weights.values())
        if total_weight != 0:
            for name in self.weights:
                self.weights[name] /= total_weight
                
        return self
        
    def build_factor_matrix(self) -> pd.DataFrame:
        """
        构建因子数据矩阵
        
        返回:
            pd.DataFrame: 因子数据矩阵
        """
        data = {}
        common_index = None
        
        for name, factor in self.factors.items():
            if factor.data is not None:
                data[name] = factor.data
                if common_index is None:
                    common_index = factor.data.index
                else:
                    common_index = common_index.intersection(factor.data.index)
                    
        if not data or common_index is None or len(common_index) == 0:
            raise ValueError("没有有效的因子数据或因子无共同索引")
            
        # 对齐索引
        aligned_data = {}
        for name, series in data.items():
            aligned_data[name] = series.loc[common_index]
            
        self.factor_data = pd.DataFrame(aligned_data)
        
        return self.factor_data
        
    def combine_factors(self, method: str = 'weighted_sum') -> pd.Series:
        """
        组合多个因子
        
        参数:
            method: 组合方法，可选 'weighted_sum', 'equal_weight', 'regression'
            
        返回:
            pd.Series: 组合因子
        """
        if self.factor_data is None:
            self.build_factor_matrix()
            
        if method == 'weighted_sum':
            # 使用指定权重
            if not self.weights:
                raise ValueError("使用weighted_sum方法需要指定因子权重")
                
            # 确保所有因子都有权重
            for name in self.factor_data.columns:
                if name not in self.weights:
                    self.weights[name] = 0
                    
            # 计算加权和
            weights_series = pd.Series({name: self.weights[name] for name in self.factor_data.columns})
            self.combined_factor = self.factor_data.mul(weights_series).sum(axis=1)
            
        elif method == 'equal_weight':
            # 等权重组合
            weight = 1.0 / len(self.factor_data.columns)
            self.weights = {name: weight for name in self.factor_data.columns}
            self.combined_factor = self.factor_data.mean(axis=1)
            
        elif method == 'regression':
            # 使用回归方法（需要目标变量）
            raise ValueError("regression方法需要使用fit_regression方法，并提供目标变量")
            
        else:
            raise ValueError(f"未知的组合方法: {method}")
            
        return self.combined_factor
        
    def fit_regression(self, target: pd.Series, model_type: str = 'linear', 
                     **kwargs) -> 'MultiFactorModel':
        """
        使用回归拟合因子权重
        
        参数:
            target: 目标变量（如收益率）
            model_type: 回归模型类型，可选 'linear', 'ridge', 'lasso'
            **kwargs: 模型参数
            
        返回:
            self: 支持链式调用
        """
        if self.factor_data is None:
            self.build_factor_matrix()
            
        # 对齐目标变量和因子数据
        common_index = self.factor_data.index.intersection(target.index)
        if len(common_index) == 0:
            raise ValueError("因子数据和目标变量无共同索引")
            
        X = self.factor_data.loc[common_index]
        y = target.loc[common_index]
        
        # 标准化特征
        scaler = StandardScaler()
        X_scaled = pd.DataFrame(
            scaler.fit_transform(X),
            index=X.index,
            columns=X.columns
        )
        
        # 创建模型
        if model_type == 'linear':
            self.model = LinearRegression(**kwargs)
        elif model_type == 'ridge':
            self.model = Ridge(**kwargs)
        elif model_type == 'lasso':
            self.model = Lasso(**kwargs)
        else:
            raise ValueError(f"未知的回归模型类型: {model_type}")
            
        # 拟合模型
        self.model.fit(X_scaled, y)
        
        # 提取权重
        self.weights = {name: coef for name, coef in zip(X.columns, self.model.coef_)}
        
        # 计算组合因子
        self.combined_factor = pd.Series(
            self.model.predict(X_scaled),
            index=X.index
        )
        
        # 为所有数据计算组合因子
        X_all_scaled = pd.DataFrame(
            scaler.transform(self.factor_data),
            index=self.factor_data.index,
            columns=self.factor_data.columns
        )
        
        self.combined_factor = pd.Series(
            self.model.predict(X_all_scaled),
            index=self.factor_data.index
        )
        
        return self
        
    def backtest(self, returns: pd.Series, n_groups: int = 5, 
                fee_rate: float = 0.0015, turnover_rate: float = 0.2) -> Dict[str, Any]:
        """
        回测组合因子
        
        参数:
            returns: 收益率数据
            n_groups: 分组数
            fee_rate: 交易费率
            turnover_rate: 换手率
            
        返回:
            Dict: 回测结果
        """
        if self.combined_factor is None:
            raise ValueError("请先使用combine_factors或fit_regression方法构建组合因子")
            
        # 创建Factor对象
        combined_factor = Factor(
            name=self.name,
            data=self.combined_factor,
            description=f"组合因子: {self.name}"
        )
        
        # 执行回测
        return combined_factor.backtest(returns, n_groups=n_groups, 
                                     fee_rate=fee_rate, turnover_rate=turnover_rate)
                                     
    def plot(self, returns: pd.Series = None, n_groups: int = 5, 
            figsize: Tuple[int, int] = (15, 12)) -> plt.Figure:
        """
        绘制多因子模型分析图表
        
        参数:
            returns: 收益率数据
            n_groups: 分组数
            figsize: 图表大小
            
        返回:
            plt.Figure: 图表对象
        """
        if self.combined_factor is None:
            raise ValueError("请先使用combine_factors或fit_regression方法构建组合因子")
            
        fig = plt.figure(figsize=figsize)
        
        # 创建布局
        if returns is not None:
            grid = plt.GridSpec(3, 2, figure=fig)
            ax1 = fig.add_subplot(grid[0, 0])  # 因子权重
            ax2 = fig.add_subplot(grid[0, 1])  # 组合因子分布
            ax3 = fig.add_subplot(grid[1, :])  # 分组收益率
            ax4 = fig.add_subplot(grid[2, :])  # 累积收益
        else:
            grid = plt.GridSpec(1, 2, figure=fig)
            ax1 = fig.add_subplot(grid[0, 0])  # 因子权重
            ax2 = fig.add_subplot(grid[0, 1])  # 组合因子分布
            ax3 = ax4 = None
            
        # 1. 因子权重
        weights_series = pd.Series(self.weights)
        weights_series = weights_series.sort_values(ascending=False)
        
        colors = ['g' if w >= 0 else 'r' for w in weights_series]
        ax1.bar(weights_series.index, weights_series, color=colors)
        ax1.set_title(f"因子权重: {self.name}")
        ax1.set_ylabel("权重")
        plt.setp(ax1.xaxis.get_majorticklabels(), rotation=45, ha='right')
        
        # 2. 组合因子分布
        sns.histplot(self.combined_factor.dropna(), ax=ax2, kde=True)
        ax2.set_title(f"组合因子分布: {self.name}")
        ax2.set_xlabel("因子值")
        ax2.set_ylabel("频率")
        
        # 如果提供了收益率数据，绘制回测相关图表
        if returns is not None and ax3 is not None and ax4 is not None:
            # 创建Factor对象
            combined_factor = Factor(
                name=self.name,
                data=self.combined_factor,
                description=f"组合因子: {self.name}"
            )
            
            backtest_result = combined_factor.backtest(returns, n_groups=n_groups)
            
            # 3. 分组收益率
            if 'annual_returns' in backtest_result:
                groups = list(backtest_result['annual_returns'].keys())
                returns_values = [backtest_result['annual_returns'][g] for g in groups]
                ax3.bar(groups, returns_values)
                ax3.set_title(f"各分组年化收益率: {self.name}")
                ax3.set_xlabel("分组")
                ax3.set_ylabel("年化收益率")
                
                # 添加数值标签
                for i, v in enumerate(returns_values):
                    ax3.text(i + 1, v, f"{v:.2%}", ha='center', va='bottom')
            
            # 4. 多空组合累积收益
            if 'cum_long_short' in backtest_result and 'cum_long_short_after_cost' in backtest_result:
                backtest_result['cum_long_short'].plot(ax=ax4, label='多空组合')
                backtest_result['cum_long_short_after_cost'].plot(ax=ax4, label='多空组合(扣费)')
                
                # 添加各分组累积收益
                if 'cum_returns' in backtest_result:
                    for group, cum_ret in backtest_result['cum_returns'].items():
                        cum_ret.plot(ax=ax4, label=f'分组{group}', alpha=0.5, linestyle='--')
                        
                ax4.set_title(f"累积收益曲线: {self.name}")
                ax4.set_xlabel("日期")
                ax4.set_ylabel("累积收益")
                ax4.legend()
                ax4.grid(True)
        
        plt.tight_layout()
        return fig
        
    def get_important_factors(self, threshold: float = 0.01) -> List[str]:
        """
        获取重要因子列表
        
        参数:
            threshold: 权重阈值
            
        返回:
            List[str]: 重要因子名称列表
        """
        return [name for name, weight in self.weights.items() if abs(weight) >= threshold]


class QFactor:
    """QTorch因子分析框架的主类"""
    
    def __init__(self):
        """初始化"""
        self.library = FactorLibrary()
        
    def create_factor(self, name: str, func: Callable = None, data: pd.Series = None, 
                    description: str = None, category: str = None) -> Factor:
        """创建因子"""
        factor = Factor(name=name, func=func, data=data, description=description, category=category)
        self.library.add_factor(factor)
        return factor
        
    def create_model(self, name: str = "多因子模型") -> MultiFactorModel:
        """创建多因子模型"""
        return MultiFactorModel(name=name)
        
    def get_factor(self, name: str) -> Optional[Factor]:
        """获取因子"""
        return self.library.get_factor(name)
        
    def list_factors(self) -> pd.DataFrame:
        """列出所有因子"""
        return self.library.list_factors()
        
    def calculate_factors(self, data: pd.DataFrame, names: List[str] = None) -> Dict[str, pd.Series]:
        """计算多个因子值"""
        return self.library.calculate_factors(data, names)
        
    def correlation_matrix(self, names: List[str] = None) -> pd.DataFrame:
        """计算因子相关性矩阵"""
        return self.library.correlation_matrix(names)
        
    def analyze_factor(self, name: str, returns: pd.Series, n_groups: int = 5) -> Dict[str, Any]:
        """分析单个因子"""
        factor = self.get_factor(name)
        if factor is None:
            raise ValueError(f"因子 {name} 不存在")
            
        # 计算IC
        ic_result = factor.calc_ic(returns)
        
        # 回测
        bt_result = factor.backtest(returns, n_groups=n_groups)
        
        return {
            'factor': factor,
            'ic_result': ic_result,
            'backtest_result': bt_result,
            'stats': factor.stats
        }
        
    def compare_factors(self, returns: pd.Series, names: List[str] = None, 
                      n_groups: int = 5) -> Dict[str, pd.DataFrame]:
        """比较多个因子"""
        ic_summary = self.library.factor_ic_summary(returns, names)
        bt_summary = self.library.backtest_summary(returns, names, n_groups)
        
        return {
            'ic_summary': ic_summary,
            'backtest_summary': bt_summary
        }
        
    def save_library(self, path: str) -> None:
        """保存因子库"""
        self.library.save(path)
        
    def load_library(self, path: str) -> None:
        """加载因子库"""
        self.library = FactorLibrary.load(path)
        
    def create_technical_factors(self, data: pd.DataFrame) -> List[Factor]:
        """创建技术分析因子"""
        factors = []
        
        # 确保数据包含必要的列
        required_columns = ['open', 'high', 'low', 'close', 'volume']
        missing_columns = [col for col in required_columns if col not in data.columns]
        if missing_columns:
            raise ValueError(f"数据缺少必要列: {missing_columns}")
        
        # 动量类因子
        factors.append(self.create_factor(
            name="momentum_1m",
            func=lambda df: df['close'].pct_change(20),
            description="1个月动量因子",
            category="动量"
        ))
        
        factors.append(self.create_factor(
            name="momentum_3m",
            func=lambda df: df['close'].pct_change(60),
            description="3个月动量因子",
            category="动量"
        ))
        
        factors.append(self.create_factor(
            name="momentum_6m",
            func=lambda df: df['close'].pct_change(120),
            description="6个月动量因子",
            category="动量"
        ))
        
        # 均线类因子
        factors.append(self.create_factor(
            name="ma_cross",
            func=lambda df: df['close'].rolling(5).mean() - df['close'].rolling(20).mean(),
            description="短期均线与长期均线差值",
            category="均线"
        ))
        
        factors.append(self.create_factor(
            name="ma_ratio",
            func=lambda df: df['close'] / df['close'].rolling(20).mean() - 1,
            description="价格与均线偏离率",
            category="均线"
        ))
        
        # 波动率类因子
        factors.append(self.create_factor(
            name="volatility_1m",
            func=lambda df: df['close'].pct_change().rolling(20).std(),
            description="1个月波动率",
            category="波动率"
        ))
        
        # RSI因子
        def calc_rsi(df, period=14):
            delta = df['close'].diff()
            gain = delta.where(delta > 0, 0)
            loss = -delta.where(delta < 0, 0)
            avg_gain = gain.rolling(window=period).mean()
            avg_loss = loss.rolling(window=period).mean()
            rs = avg_gain / avg_loss
            return 100 - (100 / (1 + rs))
            
        factors.append(self.create_factor(
            name="rsi_14",
            func=lambda df: calc_rsi(df, 14),
            description="14日RSI",
            category="超买超卖"
        ))
        
        # MACD因子
        def calc_macd(df, fast=12, slow=26, signal=9):
            ema_fast = df['close'].ewm(span=fast, adjust=False).mean()
            ema_slow = df['close'].ewm(span=slow, adjust=False).mean()
            macd_line = ema_fast - ema_slow
            signal_line = macd_line.ewm(span=signal, adjust=False).mean()
            return macd_line - signal_line
            
        factors.append(self.create_factor(
            name="macd_hist",
            func=lambda df: calc_macd(df),
            description="MACD柱状图",
            category="趋势"
        ))
        
        # 布林带因子
        def calc_bband_position(df, period=20, std_dev=2):
            middle = df['close'].rolling(window=period).mean()
            std = df['close'].rolling(window=period).std()
            upper = middle + std_dev * std
            lower = middle - std_dev * std
            return (df['close'] - lower) / (upper - lower)
            
        factors.append(self.create_factor(
            name="bb_position",
            func=lambda df: calc_bband_position(df),
            description="布林带位置",
            category="均值回归"
        ))
        
        # 成交量类因子
        factors.append(self.create_factor(
            name="volume_ratio",
            func=lambda df: df['volume'] / df['volume'].rolling(20).mean(),
            description="成交量比率",
            category="成交量"
        ))
        
        # 价格突破因子
        factors.append(self.create_factor(
            name="price_breakout",
            func=lambda df: (df['close'] - df['close'].rolling(20).max().shift(1)) / df['close'],
            description="价格突破",
            category="突破"
        ))
        
        # 计算所有因子
        for factor in factors:
            try:
                factor.calculate(data)
            except Exception as e:
                logging.error(f"计算因子 {factor.name} 时出错: {str(e)}")
                
        return factors